#!/usr/bin/env python3
# ==============================================================================
#
# Copyright (C) 2022 Sophgo Technologies Inc.  All rights reserved.
#
# TPU-MLIR is licensed under the 2-Clause BSD License except for the
# third-party components.
#
# ==============================================================================

import sys
import argparse as arg

from profile_helper.interface import *
from profile_helper.bmprofile_common import setComputeAbility


def profileArgParser():
    parser = arg.ArgumentParser()
    # yapf: disable
    parser.add_argument("--arch", type=str.upper, required=True, choices=["BM1684", "BM1684X", "BM1688", "CV186X", "BM1690", "CV184X"], help="chip arch")
    parser.add_argument("--mode", type=str, choices=["time", "command", "check", "sim", "lg"],
                        help="time: parse profile data to visualize. "
                        "sim: parse net_stat.sim in profile to visualize. "
                        "command: parse static commands in profile to analyse the command params. ",
                        default="time")
    parser.add_argument("--format", type=str, choices=["html", "csv", "console", "layer"], help="output format", default="html")
    parser.add_argument("--test", type=str, help="only for mode=command, check some condition for command", default="")
    parser.add_argument("--option", type=str, help="profile extra options, format is 'key1=value1,key2=value2,...'",
                        default="")
    parser.add_argument("--compute_ability", type=str, help="compute ability",
                        default="")
    parser.add_argument('--debug', action='store_true', help='Enable debug mode')
    parser.add_argument('--disable_doc', action='store_true', help='Disable generate doc file')
    parser.add_argument("input_dir", type=str, help="profile_data dir generated by bmruntime")
    parser.add_argument("output_dir", type=str, help="output dir that contains parsing result", default="profile_out")
    # yapf: enable
    return parser


if __name__ == "__main__":
    parser = profileArgParser()
    args, unknown = parser.parse_known_args(sys.argv[1:] + ['profile_out'])
    if args.compute_ability:
        setComputeAbility(args.compute_ability)
    if args.mode == "time":
        if args.arch in ["BM1688", "CV186X", "BM1690", "CV184X"]:
            bmprofile_parse_perfAI(args.input_dir,
                                   args.output_dir,
                                   args.test,
                                   arch=args.arch,
                                   debug=args.debug,
                                   doc=not args.disable_doc)
        else:
            bmprofile_analyze(args.input_dir, args.output_dir, args.format, args.option)
    elif args.mode == "sim":
        bmprofile_simulate(args.input_dir, args.output_dir, args.format, args.option)
    elif args.mode == "command":
        bmprofile_parse_command(args.input_dir, args.output_dir, args.test, arch=args.arch)
    elif args.mode == "check":
        bmprofile_check_command(args.input_dir, args.output_dir, args.test, arch=args.arch)
    elif args.mode == "lg":
        bmprofile_parse_layergroup_log(args.input_dir, args.output_dir, *unknown)
    else:
        assert 0, "Not support mode={}".format(args.mode)
